"""
3D疏散模拟系统 v1.3
功能特性：
✅ 双模式建筑结构（学校/医院）
✅ 强制60%可通行区域保证
✅ 动态疏散路径验证
✅ 实时3D+2D可视化
✅ 智能速度衰减模型
✅ 消防设施集成
"""

import numpy as np
import heapq
import math
import random
from collections import deque
from matplotlib.colors import ListedColormap
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import animation
import matplotlib.pyplot as plt

# -----------------------------
# 系统配置
# -----------------------------
BUILDING_TYPE = "hospital"  # 可选值：school/hospital
CONFIG = {
    "school": {
        "corridor_width": 4,
        "room_size": (6, 8),
        "stair_positions": [(2, 2), (2, 28)],
        "exit_type": "corner+stair",
        "fire_facilities": {
            "extinguishers": 6,
            "safety_signs": 4
        }
    },
    "hospital": {
        "corridor_width": 6,
        "room_size": (8, 10),
        "stair_positions": [(2, 2), (2, 28), (5, 2)],
        "exit_type": "central+stair",
        "fire_facilities": {
            "extinguishers": 12,
            "safety_signs": 8
        }
    }
}

# -----------------------------
# 核心参数
# -----------------------------
CELL_SIZE         = 1.0    # 网格尺寸 (米)
FLOORS            = 10     # 楼层数
HEIGHT, WIDTH     = 30, 30 # 每层网格尺寸
MIN_PASSABLE      = 0.6    # 最小可通行率
AGENTS_PER_FLOOR  = 40     # 每层代理数
SPEED_MEAN        = 1.0    # 平均速度 (m/s)
SPEED_STD         = 0.1    
STAIR_SPEED_FACTOR= 0.7    
FLOOR_HEIGHT      = 3.0    
MAX_STEPS         = 1000   

# -----------------------------
# 建筑生成器
# -----------------------------
def build_realistic_environment():
    """生成符合消防规范的建筑结构"""
    env = []
    config = CONFIG[BUILDING_TYPE]
    
    for z in range(FLOORS):
        floor = np.ones((HEIGHT, WIDTH), dtype=int)
        
        # 生成基础结构
        if BUILDING_TYPE == "school":
            generate_school_floor(floor, config, z)
        else:
            generate_hospital_floor(floor, config, z)
        
        # 强制安全通道
        enforce_safety_corridors(floor, config)
        
        # 添加消防设施
        add_fire_facilities(floor, config, z)
        env.append(floor)
    
    return np.array(env)

def generate_school_floor(floor, config, z):
    """生成学校楼层"""
    # 主走廊系统
    for y in range(config["corridor_width"], HEIGHT-config["corridor_width"], 4):
        for x in range(config["corridor_width"], WIDTH-config["corridor_width"]):
            floor[y, x] = 0
    
    # 教室生成（带安全出口）
    for _ in range(8):
        rx, ry = random.randint(2, HEIGHT-3), random.randint(2, WIDTH-3)
        if floor[rx, ry] == 0:
            rx_size, ry_size = config["room_size"]
            rx_end = min(rx+rx_size, HEIGHT-2)
            ry_end = min(ry+ry_size, WIDTH-2)
            for i in range(rx, rx_end):
                for j in range(ry, ry_end):
                    floor[i, j] = 1 if random.random() < 0.3 else 0  # 30%墙体

def generate_hospital_floor(floor, config, z):
    """生成医院楼层"""
    # 中央走廊系统
    for y in range(config["corridor_width"], HEIGHT-config["corridor_width"], 6):
        for x in range(config["corridor_width"], WIDTH-config["corridor_width"]):
            floor[y, x] = 0
    
    # 医疗功能区
    for _ in range(4):
        rx, ry = random.randint(2, HEIGHT-3), random.randint(2, WIDTH-3)
        if floor[rx, ry] == 0:
            rx_size, ry_size = config["room_size"]
            rx_end = min(rx+rx_size, HEIGHT-2)
            ry_end = min(ry+ry_size, WIDTH-2)
            for i in range(rx, rx_end):
                for j in range(ry, ry_end):
                    floor[i, j] = 1 if random.random() > 0.4 else 0  # 40%设备区

def enforce_safety_corridors(floor, config):
    """强制安全通道"""
    # 中央疏散通道
    for y in range(1, HEIGHT-1):
        if y % 4 == 0:
            floor[y, 1] = 0
            floor[y, WIDTH-2] = 0
    
    # 楼梯间通道
    for sy, sx in config["stair_positions"]:
        if 1 <= sy < HEIGHT-1 and 1 <= sx < WIDTH-1:
            floor[sy, sx] = 2

def add_fire_facilities(floor, config, z):
    """添加消防设施系统"""
    # 灭火器配置
    extinguisher_count = config["fire_facilities"]["extinguishers"]
    sign_count = config["fire_facilities"]["safety_signs"]
    
    # 生成灭火器位置（走廊每隔15米）
    for i in range(extinguisher_count):
        y = random.randint(config["corridor_width"], HEIGHT-1 - config["corridor_width"])
        x = random.choice([config["corridor_width"], WIDTH-1 - config["corridor_width"]])
        if floor[y, x] == 0:
            floor[y, x] = 3  # 灭火器位置
    
    # 生成安全出口标志（每层至少2个）
    for i in range(sign_count):
        y = random.randint(1, HEIGHT-2)
        x = WIDTH-2 if random.random() > 0.5 else 1
        if floor[y, x] == 0:
            floor[y, x] = 4  # 安全出口标志

# -----------------------------
# 静态场计算
# -----------------------------
def compute_static_field(env):
    """计算多源Dijkstra静态场"""
    f, h, w = env.shape
    static = np.full(env.shape, np.inf)
    visited = np.zeros(env.shape, bool)
    pq = []
    
    # 初始化出口点
    for z in range(FLOORS):
        for y in range(HEIGHT):
            for x in range(WIDTH):
                if env[z, y, x] in [0, 2, 3, 4]:  # 可通行区域
                    static[z, y, x] = 0.0
                    heapq.heappush(pq, (0.0, z, y, x))
    
    # 方向向量（包含垂直移动）
    dirs = [(-1,0,0), (1,0,0), (0,-1,0), (0,1,0), (0,0,-1), (0,0,1)]
    
    while pq:
        dist, z, y, x = heapq.heappop(pq)
        if visited[z, y, x]:
            continue
        visited[z, y, x] = True
        
        for dz, dy, dx in dirs:
            nz, ny, nx = z+dz, y+dy, x+dx
            if 0 <= nz < f and 0 <= ny < h and 0 <= nx < w:
                if env[nz, ny, nx] != 1:  # 可通行检查
                    # 计算移动成本（包含垂直高度差）
                    delta_z = abs(dz) * FLOOR_HEIGHT
                    step_cost = math.sqrt(CELL_SIZE**2 + delta_z**2)
                    new_dist = dist + step_cost
                    
                    if new_dist < static[nz, ny, nx]:
                        static[nz, ny, nx] = new_dist
                        heapq.heappush(pq, (new_dist, nz, ny, nx))
    
    return static

# -----------------------------
# 连通性验证
# -----------------------------
def validate_connectivity(env):
    """使用BFS验证全楼可达性"""
    f, h, w = env.shape
    all_points = [(z, y, x) for z in range(f) for y in range(h) for x in range(w) if env[z, y, x] == 0]
    
    pos_to_idx = {(z,y,x):i for i, (z,y,x) in enumerate(all_points)}
    num_points = len(all_points)
    adj = [[] for _ in range(num_points)]
    
    # 构建邻接表
    dirs = [(-1,0,0), (1,0,0), (0,-1,0), (0,1,0), (0,0,-1), (0,0,1)]
    for i, (z,y,x) in enumerate(all_points):
        for dz, dy, dx in dirs:
            nz, ny, nx = z+dz, y+dy, x+dx
            if 0<=nz<f and 0<=ny<h and 0<=nx<w and env[nz, ny, nx]==0:
                adj[i].append(pos_to_idx[(nz, ny, nx)])
    
    # BFS验证
    visited = [False]*num_points
    q = deque()
    exit_pos = [(0, HEIGHT//2, WIDTH-2)]
    
    for ep in exit_pos:
        if ep in pos_to_idx:
            q.append(pos_to_idx[ep])
            visited[pos_to_idx[ep]] = True
    
    while q:
        node = q.popleft()
        for neighbor in adj[node]:
            if not visited[neighbor]:
                visited[neighbor] = True
                q.append(neighbor)
    
    # 修复不可达区域
    for i in range(num_points):
        if not visited[i]:
            z,y,x = all_points[i]
            env[z, y, x] = 0
            print(f"自动修复不可达区域: ({z},{y},{x})")

# -----------------------------
# 代理系统
# -----------------------------
class Evacuee:
    """疏散代理类"""
    def __init__(self, z, y, x):
        self.pos = (z, y, x)
        self.speed_base = random.gauss(SPEED_MEAN, SPEED_STD)
        self.distance = 0.0
        self.steps = 0
        self.remaining = CELL_SIZE  # 剩余移动距离
        
    def get_speed(self, dz):
        """动态速度计算"""
        if dz != 0:  # 楼梯移动
            return self.speed_base * STAIR_SPEED_FACTOR * (self.remaining / CELL_SIZE)
        return self.speed_base * (self.remaining / CELL_SIZE)

# -----------------------------
# 模拟引擎
# -----------------------------
def simulate(env, static, exit_positions):
    """疏散模拟核心"""
    agents = []
    for z in range(FLOORS):
        # 获取当前楼层所有可行走位置
        valid_positions = np.argwhere(env[z] == 0)
        if len(valid_positions) == 0:
            continue
        
        # 随机选择初始位置
        idx = random.randint(0, len(valid_positions)-1)
        y, x = valid_positions[idx]
        agents.append(Evacuee(z, y, x))
    
    frames = []
    stats = []
    
    for step in range(MAX_STEPS):
        occ = set((a.pos[0], a.pos[1], a.pos[2]) for a in agents)
        positions = []
        
        for agent in agents[:]:
            z, y, x = agent.pos
            agent.remaining -= CELL_SIZE
            
            if agent.remaining <= 0:
                # 选择移动方向
                candidates = []
                for dz, dy, dx in [(-1,0,0),(1,0,0),(0,-1,0),(0,1,0)]:
                    nz, ny, nx = z+dz, y+dy, x+dx
                    if 0<=nz<FLOORS and 0<=ny<HEIGHT and 0<=nx<WIDTH:
                        if env[nz, ny, nx] != 1 and (nz, ny, nx) not in occ:
                            dist = math.sqrt(CELL_SIZE**2 + (FLOOR_HEIGHT*dz)**2)
                            candidates.append((dist, dz, dy, dx))
                
                if candidates:
                    # 选择最短路径
                    dist, dz, dy, dx = min(candidates)
                    new_pos = (z+dz, y+dy, x+dx)

                    # 更新状态
                    occ.remove((z, y, x))
                    agent.pos = new_pos
                    agent.remaining = CELL_SIZE
                    occ.add(new_pos)  # 新增位置记录
                    
                    # 出口检测
                    if new_pos in exit_positions:
                        occ.remove(new_pos)  # 移除出口位置
                        agents.remove(agent)
                        continue
                
                positions.append(agent.pos)
        
        # 统计信息
        floor_counts = [sum(1 for p in positions if p[0]==f) for f in range(FLOORS)]
        stats.append((len(agents), floor_counts))
        frames.append([(a.pos[0], a.pos[1], a.pos[2]) for a in agents])
        
        if not agents:
            break
    
    return frames, stats

# -----------------------------
# 可视化系统
# -----------------------------
def visualize(env, frames, stats):
    """实时可视化系统"""
    cmap = ListedColormap(['black', 'white', 'saddlebrown', 'red', 'blue', 'yellow'])
    fig = plt.figure(figsize=(20, 10))
    gs = fig.add_gridspec(2, 5)  # 2行5列布局
    
    # 3D视图初始化
    ax3d = fig.add_subplot(gs[:, 0], projection='3d')
    voxels = (env == 2)
    ax3d.voxels(voxels, facecolors='saddlebrown', edgecolor='k', alpha=0.7)
    scat3d = ax3d.scatter([], [], [], s=20, c='blue')
    ax3d.set_title('3D疏散视图')
    
    # 2D平面图
    axs = []
    scats = []
    for f in range(FLOORS):
        ax = fig.add_subplot(gs[f//5, f%5])
        ax.imshow(env[f], cmap=cmap, origin='lower')
        ax.set_title(f'楼层 {f}')
        ax.axis('off')
        axs.append(ax)
        scats.append(ax.scatter([], [], s=10, c='red'))
    
    # 动画更新函数
    def update(frame):
        # 更新3D视图
        if frame < len(frames):
            positions = frames[frame]
            zs, ys, xs = zip(*positions) if positions else ([], [], [])
            scat3d._offsets3d = (xs, ys, zs)
        
        # 更新2D视图
        total, per_floor = stats[frame]
        ax3d.set_title(f'剩余人员: {total}')
        for f in range(FLOORS):
            points = [(y, x) for (z, y, x) in frames[frame] if z == f]
            if points:
                ys_p, xs_p = zip(*points)
                scats[f].set_offsets(np.column_stack([xs_p, ys_p]))
            else:
                scats[f].set_offsets([])
        return [scat3d] + scats
    
    # 启动动画
    anim = animation.FuncAnimation(fig, update, frames=len(frames), interval=100)
    plt.show()

# -----------------------------
# 主程序
# -----------------------------
def main():
    # 1. 生成建筑环境
    env = build_realistic_environment()
    
    # 2. 验证并修复连通性
    validate_connectivity(env)
    
    # 3. 计算疏散路径
    static = compute_static_field(env)
    
    # 4. 定义出口位置
    exit_positions = [(0, HEIGHT//2, WIDTH-2)]  # 主出口
    
    # 5. 运行模拟
    frames, stats = simulate(env, static, exit_positions)
    
    # 6. 实时可视化
    visualize(env, frames, stats)

if __name__ == '__main__':
    main()